#include "monitor/Monitor.h"
#include "types.h"

#include <cstddef>
#include <cstdlib>
#include <exception>
#include <getopt.h>
#include <readline/history.h>
#include <readline/readline.h>

static const uint32_t dimg[] = {
    0x00000297, // auipc t0,0
    0x0002b823, // sd  zero,16(t0)
    0x0102b503, // ld  a0,16(t0)
    0x00100073, // ebreak (used as nemu_trap)
    0xdeadbeef, // some data
};

// static const uint32_t dimg[] = {
//     0x000085b7,
//     0x4215859b,
//     0x01059593,
//     0x01059593,
//     0x01059593,
//     0x000f1637,
//     0x23f6061b,
//     0x02c586b3,
//     0x02c59733,
//     0x02c5a7b3,
//     0x02c5b833,
//     0xfff0059b,
//     0x00058613,
//     0x02c586b3,
//     0x02c59733,
//     0x02c5a7b3,
//     0x02c5b833,
//     0x00100073,
// };

int Monitor::cmd_s() {
    int n = 1;
    std::string arg;
    if (iss >> arg) n = std::atoi(arg.data());
    this->run(n);
    state = State::PAUSE;
    return 0;
}

int Monitor::cmd_c() {
    this->run(-1);
    state = State::PAUSE;
    return 0;
}

int Monitor::cmd_e() {
    std::string arg;

    if (!std::getline(iss, arg)) return 0;
    
    auto [success, val] = expr->parser(arg.c_str());
    if (!success) {
        std::cout << "Bad Expression: " << arg << std::endl;
    } else {
        std::cout << arg << " = ";
        std::cout << "[" << std::dec << val << ", " << std::showbase << std::hex << val << "]" << std::endl;
    }
    return 0;
}

int Monitor::cmd_x() {
    std::string arg;
    if (!(iss >> arg)) return 0;

    int N = std::atoi(arg.data());

    if (!(iss >> arg)) return 0;
    auto [success, base_addr] = expr->parser(arg.c_str());
    if (!success) {
        std::cout << "Bad Expression: " << arg << std::endl;
        return 0;
    }
    int word_size = sizeof(word_t);
    for (int i = 0; i < N; ++i) {
        word_t data;
        ram->read(base_addr + i * word_size, word_size, (uint8_t*)&data);
        printf(FMT_WORD ": ", (base_addr + i * word_size));
        for (int j = 0; j < word_size; ++j) {
            printf("0x%02x  ", (uint8_t) (data&0xFF));
            data >>= word_size;
        }
        printf("\n");
    }
    return 0;
}

int Monitor::cmd_i() {
#ifdef CONFIG_ITRACE
    std::string arg;
    iss >> arg;
    if (arg == "on") {
        core->itrace->enable = true;
        return 0;
    }
    if (arg == "off") {
        core->itrace->enable = false;
        return 0;
    }
    std::printf("inst count: %ld\n", core->inst_count);
    std::printf("cycle count: %ld\n", core->cycle_count);
    std::printf("last pc: " FMT_WORD "\n", core->last_pc);
    std::printf("last inst: %#x\n", core->last_inst);
    core->itrace->print();
#endif
    return 0;
}

int Monitor::cmd_r() {
    core->dump_regs();
    return 0;
}

int Monitor::cmd_v() {
#ifdef CONFIG_VTRACE
    std::string arg;
    if (!(iss >> arg)) return 0;

    if (arg == "on") {
        core->vtrace->enable = true;
        return 0;
    }
    if (arg == "off") {
        core->vtrace->enable = false;
        return 0;
    }
#endif
    return 0;
}

int Monitor::cmd_w() {
#ifdef CONFIG_WATCHPOINT
    std::string arg;
    if (!(iss >> arg)) return 0;

    if (arg == "add") {
        if (!(iss >> arg)) return 0;
        word_t wp = std::strtoul(arg.c_str(), NULL, 0);
        core->watchpoint.add(wp);
        return 0;
    }
    if (arg == "del") {
        if (!(iss >> arg)) return 0;
        word_t wp = std::strtoul(arg.c_str(), NULL, 0);
        core->watchpoint.del(wp);
        return 0;
    }
    if (arg == "list") {
        core->watchpoint.list();
        return 0;
    }
    if (arg == "clear") {
        core->watchpoint.clear();
        return 0;
    }
#endif
    return 0;
}

int Monitor::cmd_q() {
    state = State::QUIT;
    return -1;
}

void Monitor::cmd_table_init() {
    expr = new Expression([&](std::string& name, bool& success) -> word_t {
        try {
            if (name == "$pc") return core->pc;
            return core->gpr[name.substr(1)];
        } catch (std::runtime_error& e) {
            success = false;
            return 0;
        }
    });

    cmd_table.push_back({"s", "step n instructions", std::bind(&Monitor::cmd_s, this)});
    cmd_table.push_back({"c", "continue execution", std::bind(&Monitor::cmd_c, this)});
    cmd_table.push_back({"e", "evaluate expression", std::bind(&Monitor::cmd_e, this)});
    cmd_table.push_back({"x", "examine memory", std::bind(&Monitor::cmd_x, this)});
    cmd_table.push_back({"i", "print instruction trace", std::bind(&Monitor::cmd_i, this)});
    cmd_table.push_back({"r", "print register information", std::bind(&Monitor::cmd_r, this)});
    cmd_table.push_back({"v", "print instruction trace", std::bind(&Monitor::cmd_v, this)});
    cmd_table.push_back({"w", "watch points", std::bind(&Monitor::cmd_w, this)});
    cmd_table.push_back({"q", "quit", std::bind(&Monitor::cmd_q, this)});
}

void Monitor::device_init(const char* img_file) {
    
    if (img_file != NULL && std::filesystem::exists(img_file)) { 
        Log("The image is %s", img_file);
        ram = new Memory(img_file); 
    } else {
        Log("No image is given. Use the default build-in image.\n");
        ram = new Memory((uint8_t*)dimg, sizeof(dimg)); 
    }
    xbar.add_device(MEM_BASE, ram->get_size(), ram);
    serial = new Serial();
    xbar.add_device(SERIAL_BASE, SERIAL_SIZE, serial);

    timer = new Timer();
    xbar.add_device(TIMER_BASE, TIMER_SIZE, timer);

#if CONFIG_CLINT
    clint = new CLINT();
    xbar.add_device(CLINT_BASE, CLINT_SIZE, clint);
#endif

#if CONFIG_KEYBOARD
    keyboard = new Keyboard();
    xbar.add_device(KEYBOARD_BASE, KEYBOARD_SIZE, keyboard);
#endif

#if CONFIG_DISPLAY
    display = new Display();
    xbar.add_device(VCTL_BASE, VCTL_SIZE, display);
    xbar.add_device(VMEM_BASE, VMEM_SIZE, display);
#endif

    device_thread = new std::thread([&]() -> void {
        while (state != State::QUIT) {
            xbar.do_update();
            std::this_thread::sleep_for(std::chrono::milliseconds(1000 / 60));
        }
    });
}

Monitor::Monitor(int argc, char *argv[])
{
    char* img_file = NULL;
    char* ref_file = NULL;
    char* log_path = NULL;
    bool is_batch_mode = false;

    const struct option table[] = {
        {"batch"    , no_argument      , NULL, 'b'},
        {"img"      , required_argument, NULL, 'i'},
        {"ref"      , required_argument, NULL, 'r'},
        {"log"      , required_argument, NULL, 'l'},
        {"help"     , no_argument      , NULL, 'h'},
        {0          , 0                , NULL,  0 },
    };
    int o;
    while ( (o = getopt_long(argc, argv, "-bhi:r:", table, NULL)) != -1) {
        switch (o) {
            case 'b': is_batch_mode = true; break;
            case 'i': img_file = optarg; break;
            case 'r': ref_file = optarg; break;
            case 'l': log_path = optarg; break;
            default:
                printf("Usage: %s [OPTION...] IMAGE [args]\n\n", argv[0]);
                printf("\t-b,--batch              run with batch mode\n");
                printf("\t-l,--log=FILE           output log to FILE\n");
                printf("\n");
                exit(0);
        }
    }

    this->batch_mode = is_batch_mode;

    device_init(img_file);

    cmd_table_init();
    
    core = new Core();
    core->set_xbar(&xbar);
    core->open_trace(log_path);
    core->reset(5);

#ifdef CONFIG_DETECT
    detect_thread = new std::thread([&]() -> void {
        size_t last_inst_count = 0;
        while (state != State::QUIT) {
            last_inst_count = core->inst_count;
            std::this_thread::sleep_for(std::chrono::milliseconds(1000));
            if (state != State::RUN) continue;
            if (last_inst_count == core->inst_count) {
                std::printf(ANSI_FMT("DETECT Dead Lock", ANSI_FG_RED) " at pc = " FMT_WORD 
                            ", executed instruction = %ld" "\n", core->pc, core->inst_count);
                state = State::PAUSE;
                break;
            }
        }
    });
#endif

    IFDEF(CONFIG_DIFFTEST, core->difftest_init(ref_file, ram->get_base(), ram->get_size()));
}

Monitor::~Monitor()
{
    if (device_thread != nullptr && device_thread->joinable()) { 
        device_thread->join(); 
        delete device_thread; 
    };
#ifdef CONFIG_DETECT
    if (detect_thread != nullptr && detect_thread->joinable()) { 
        detect_thread->join();
        delete detect_thread;
    };
#endif
    if (core != nullptr) delete core;
    if (clint != nullptr) delete clint;
    if (ram != nullptr) delete ram;
    if (expr != nullptr) delete expr;
    if (serial != nullptr) delete serial;
    if (timer != nullptr) delete timer;
    if (keyboard != nullptr) delete keyboard;
    if (display != nullptr) delete display;
}

void Monitor::run(size_t n) {
    if (state == State::END || state == State::ABORT) {
        std::cout << "Program execution has ended. To restart the program, exit NEMU and run again." << std::endl;
        return;
    }
    bool msi = false, mti = false;
    while (n --) {
        try {
#ifdef CONFIG_CLINT
            clint->tick();
            msi = clint->msi(), mti = clint->mti();
#endif
            state = State::RUN;
            core->step(msi, mti);
#ifdef CONFIG_WATCHPOINT
            if (core->stop_flag) {
                std::printf(ANSI_FMT("PAUSE", ANSI_FG_YELLOW) " at pc = " FMT_WORD "\n", core->pc);
                state = State::PAUSE;
                break; 
            }
#endif
        } catch (RVException& e) {
            if (e.get_trap_cause() == Trap::Breakpoint) {
                state = State::END;
                halt_value = e.get_trap_val();
                if (halt_value == 0) {
                    std::printf(ANSI_FMT("HIT GOOD TRAP", ANSI_FG_GREEN) "\n");
                } else {
                    std::printf(ANSI_FMT("HIT BAD TRAP", ANSI_FG_RED) "\n");
                }
                size_t icache_miss = core->top->icache_miss_count;
                size_t dcache_miss = core->top->dcache_miss_count;
                size_t mmio_inst = core->top->mmio_inst_count;
                size_t inst_mem = core->top->mem_inst_count;
                size_t commit_inst = core->top->commit_inst_count;
                size_t cycle_count = core->cycle_count;

                std::printf(ANSI_FMT("Statistics: ", ANSI_FG_YELLOW) "\n");
                std::printf("\t" ANSI_FMT("icache_miss_count: ", ANSI_FG_CYAN) "%ld\n", icache_miss);
                std::printf("\t" ANSI_FMT("dcache_miss_count: ", ANSI_FG_CYAN) "%ld\n", dcache_miss);
                std::printf("\t" ANSI_FMT("mmio_inst_count: ", ANSI_FG_CYAN) "%ld\n", mmio_inst);
                std::printf("\t" ANSI_FMT("mem_inst_count: ", ANSI_FG_CYAN) "%ld\n", inst_mem);
                std::printf("\t" ANSI_FMT("commit_inst_count: ", ANSI_FG_CYAN) "%ld\n", commit_inst);
                std::printf("\t" ANSI_FMT("cycle_count: ", ANSI_FG_CYAN) "%ld\n", cycle_count);
                std::printf("\t" ANSI_FMT("ICache Hit Rate: ", ANSI_FG_CYAN) "%f\n", 1.0 * (commit_inst - icache_miss) / commit_inst);
                std::printf("\t" ANSI_FMT("DCache Hit Rate: ", ANSI_FG_CYAN) "%f\n", 1.0 * (commit_inst - dcache_miss) / commit_inst);
                std::printf("\t" ANSI_FMT("IPC: ", ANSI_FG_CYAN) "%f\n", 1.0 * commit_inst / cycle_count);
            } else {
                std::cout << e.what() << std::endl;
                state = State::ABORT;
                std::printf(ANSI_FMT("ABORT", ANSI_FG_RED) " at pc = " FMT_WORD "\n", core->pc);
            }
            return;
        }
    }
}

int Monitor::loop() {
    if (batch_mode) {
        this->run(-1);
        return !(state == State::END && halt_value == 0 || state == State::QUIT);
    }
    std::string line, cmd;
    std::cout << "Welcome to npc!" << std::endl;
    using_history();

    while (true) {
        char* s = readline("(npc) ");
        if (s && *s) { add_history(s); line = s; }
        if (s != NULL) { free(s); }
        else { continue; }

        iss.clear();
        iss.str(line); 
        iss >> cmd;
        
        auto cmd_it = std::find_if(cmd_table.begin(), cmd_table.end(), [&](Command& c) -> bool {
            return c.name == cmd;
        });
        if (cmd_it == cmd_table.end()) {
            std::cout << "Unknown command: " << cmd << std::endl;
            continue;
        }
        if (cmd_it->handler() < 0) { break; }
    }
    return !(state == State::END && halt_value == 0 || state == State::QUIT);
}
